package com.ecar.oauth.provider.beans;

import com.ecar.cache.RedisFactory;
import org.springframework.security.oauth2.common.ExpiringOAuth2RefreshToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2RefreshToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.AuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.DefaultAuthenticationKeyGenerator;
import org.springframework.security.oauth2.provider.token.TokenStore;
import org.springframework.security.oauth2.provider.token.store.redis.JdkSerializationStrategy;
import org.springframework.security.oauth2.provider.token.store.redis.RedisTokenStoreSerializationStrategy;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.Pipeline;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.List;

/**
 * Created by 30 on 2017/1/25.
 * Duplication of com.wisme.sky.api.oauth.OauthRedisStore
 */
public class OAuthRedisStore implements TokenStore {

	/** +accessToken.getValue() 存储authentication对象 */
	private static final String AUTH = "auth:";
	/** +accessToken.getValue() 存储accessToken对象 */
	private static final String ACCESS = "access:";
	/** +serialize(authentication) 存储accessToken对象 */
	private static final String AUTH_TO_ACCESS = "auth_to_access:";
	/** +refreshToken.getValue() 存储refreshToken对象 */
	private static final String REFRESH = "refresh:";
	/** +accessToken.getValue() 存储refreshToken.getValue() */
	private static final String ACCESS_TO_REFRESH = "access_to_refresh:";
	/** +refreshToken.getValue() 存储accessToken.getValue() */
	private static final String REFRESH_TO_ACCESS = "refresh_to_access:";
	/** +refreshToken.getValue() 存储authentication对象 */
	private static final String REFRESH_TO_AUTH = "refresh_to_auth:";
	/** +clientId 存储accessToken对象列表 */
	private static final String CLIENT_ID_TO_ACCESS = "client_to_access:";
	/** +clientId+username 存储accessToken对象列表 */
	private static final String UNAME_TO_ACCESS = "uname_to_access:";

	private final RedisFactory conn = RedisFactory.getInstance();
	private AuthenticationKeyGenerator authenticationKeyGenerator = new DefaultAuthenticationKeyGenerator();
	private RedisTokenStoreSerializationStrategy serializationStrategy = new JdkSerializationStrategy();
	private String prefix = "oauth:";

	public void setAuthenticationKeyGenerator(AuthenticationKeyGenerator authenticationKeyGenerator) {
		this.authenticationKeyGenerator = authenticationKeyGenerator;
	}

	public void setSerializationStrategy(RedisTokenStoreSerializationStrategy serializationStrategy) {
		this.serializationStrategy = serializationStrategy;
	}

	public void setPrefix(String prefix) {
		this.prefix = prefix;
	}

	private byte[] serialize(Object object) {
		return serializationStrategy.serialize(object);
	}

	private byte[] serializeKey(String object) {
		return serialize(this.prefix + object);
	}

	private OAuth2AccessToken deserializeAccessToken(byte[] bytes) {
		return serializationStrategy.deserialize(bytes, OAuth2AccessToken.class);
	}

	private OAuth2Authentication deserializeAuthentication(byte[] bytes) {
		return serializationStrategy.deserialize(bytes, OAuth2Authentication.class);
	}

	private OAuth2RefreshToken deserializeRefreshToken(byte[] bytes) {
		return serializationStrategy.deserialize(bytes, OAuth2RefreshToken.class);
	}

	private byte[] serialize(String string) {
		return serializationStrategy.serialize(string);
	}

	private String deserializeString(byte[] bytes) {
		return serializationStrategy.deserializeString(bytes);
	}

	public OAuth2AccessToken getAccessToken(OAuth2Authentication authentication) {
		String key = authenticationKeyGenerator.extractKey(authentication);
		byte[] bytes = conn.get(serializeKey(AUTH_TO_ACCESS + key));

		OAuth2AccessToken accessToken = deserializeAccessToken(bytes);
		if (accessToken != null) {
			String extractedKey = authenticationKeyGenerator.extractKey(readAuthentication(accessToken.getValue()));
			if (!key.equals(extractedKey)) storeAccessToken(accessToken, authentication);
		}
		return accessToken;
	}

	public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
		return readAuthentication(token.getValue());
	}

	public OAuth2Authentication readAuthentication(String token) {
		byte[] bytes = conn.get(serializeKey(AUTH + token));
		return deserializeAuthentication(bytes);
	}

	public OAuth2Authentication readAuthenticationForRefreshToken(OAuth2RefreshToken token) {
		return readAuthenticationForRefreshToken(token.getValue());
	}

	public OAuth2Authentication readAuthenticationForRefreshToken(String tokenValue) {
		byte[] bytes = conn.get(serializeKey(REFRESH_TO_AUTH + tokenValue));
		return deserializeAuthentication(bytes);
	}

	public void storeAccessToken(OAuth2AccessToken accessToken, OAuth2Authentication authentication) {
		OAuth2RefreshToken refreshToken = accessToken.getRefreshToken();
		String accessTokenValue = accessToken.getValue();
		String refreshTokenValue = refreshToken.getValue();

		byte[] serializedAccessToken = serialize(accessToken);
		byte[] serializedAuth = serialize(authentication);
		byte[] accessKey = serializeKey(ACCESS + accessTokenValue);
		byte[] authKey = serializeKey(AUTH + accessTokenValue);
		byte[] authToAccessKey = serializeKey(AUTH_TO_ACCESS + authenticationKeyGenerator.extractKey(authentication));
		byte[] approvalKey = serializeKey(UNAME_TO_ACCESS + getApprovalKey(authentication));
		byte[] clientId = serializeKey(CLIENT_ID_TO_ACCESS + authentication.getOAuth2Request().getClientId());

		try (Jedis jedis = conn.getResource()) {
			Pipeline pip = jedis.pipelined();
			pip.set(accessKey, serializedAccessToken);
			pip.set(authKey, serializedAuth);
			pip.set(authToAccessKey, serializedAccessToken);
			if (!authentication.isClientOnly()) {
				pip.lrem(approvalKey, 1L, serializedAccessToken);
				pip.rpush(approvalKey, serializedAccessToken);
			}

			pip.lrem(clientId, 1L, serializedAccessToken);
			pip.rpush(clientId, serializedAccessToken);
			if (accessToken.getExpiration() != null) {
				int expiresIn = accessToken.getExpiresIn();
				pip.expire(accessKey, expiresIn);
				pip.expire(authKey, expiresIn);
				pip.expire(authToAccessKey, expiresIn);
				pip.expire(clientId, expiresIn);
				pip.expire(approvalKey, expiresIn);
			}

			if (refreshToken != null && refreshTokenValue != null) {
				byte[] refresh = serialize(refreshTokenValue);
				byte[] access = serialize(accessTokenValue);
				byte[] refreshToAccessKey = serializeKey(REFRESH_TO_ACCESS + refreshTokenValue);
				byte[] accessToRefreshKey = serializeKey(ACCESS_TO_REFRESH + accessTokenValue);
				pip.set(refreshToAccessKey, access);
				pip.set(accessToRefreshKey, refresh);
				if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
					ExpiringOAuth2RefreshToken expiringRefreshToken = (ExpiringOAuth2RefreshToken) refreshToken;
					Date expiration = expiringRefreshToken.getExpiration();
					if (expiration != null) {
						int seconds = Long.valueOf((expiration.getTime() - System.currentTimeMillis()) / 1000L).intValue();
						pip.expire(refreshToAccessKey, seconds);
						pip.expire(accessToRefreshKey, seconds);
					}
				}
			}

			pip.sync();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	private static String getApprovalKey(OAuth2Authentication authentication) {
		String userName = "";
		if (authentication.getUserAuthentication() != null)
			userName = authentication.getUserAuthentication().getName();

		return getApprovalKey(authentication.getOAuth2Request().getClientId(), userName);
	}

	private static String getApprovalKey(String clientId, String userName) {
		return clientId + (userName == null ? "" : ":" + userName);
	}

	public void removeAccessToken(OAuth2AccessToken accessToken) {
		removeAccessToken(accessToken.getValue());
	}

	public OAuth2AccessToken readAccessToken(String tokenValue) {
		byte[] bytes = conn.get(serializeKey(ACCESS + tokenValue));
		return deserializeAccessToken(bytes);
	}

	public void removeAccessToken(String tokenValue) {
		byte[] accessKey = serializeKey(ACCESS + tokenValue);
		byte[] authKey = serializeKey(AUTH + tokenValue);
		byte[] accessToRefreshKey = serializeKey(ACCESS_TO_REFRESH + tokenValue);

		try (Jedis jedis = conn.getResource()) {
			Pipeline pip = jedis.pipelined();
			pip.get(accessKey);
			pip.get(authKey);
			pip.del(accessKey);
			pip.del(accessToRefreshKey);
			pip.del(authKey);
			List results = pip.syncAndReturnAll();

			byte[] access = (byte[]) results.get(0);
			byte[] auth = (byte[]) results.get(1);
			OAuth2Authentication authentication = deserializeAuthentication(auth);
			if (authentication != null) {
				String extractedKey = authenticationKeyGenerator.extractKey(authentication);
				byte[] authToAccessKey = serializeKey(AUTH_TO_ACCESS + extractedKey);
				byte[] unameKey = serializeKey(UNAME_TO_ACCESS + getApprovalKey(authentication));
				byte[] clientId = serializeKey(CLIENT_ID_TO_ACCESS + authentication.getOAuth2Request().getClientId());

				Pipeline pip2 = jedis.pipelined();
				pip2.del(authToAccessKey);
				pip2.lrem(unameKey, 1L, access);
				pip2.lrem(clientId, 1L, access);
				pip2.del(serialize(ACCESS + extractedKey));
				pip2.sync();
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public void storeRefreshToken(OAuth2RefreshToken refreshToken, OAuth2Authentication authentication) {
		String refreshTokenValue = refreshToken.getValue();
		byte[] refreshKey = serializeKey(REFRESH + refreshTokenValue);
		byte[] refreshAuthKey = serializeKey(REFRESH_TO_AUTH + refreshTokenValue);
		byte[] serializedRefreshToken = serialize(refreshToken);

		try (Jedis jedis = conn.getResource()) {
			Pipeline pip = jedis.pipelined();
			pip.set(refreshKey, serializedRefreshToken);
			pip.set(refreshAuthKey, serialize(authentication));
			if (refreshToken instanceof ExpiringOAuth2RefreshToken) {
				ExpiringOAuth2RefreshToken expiringRefreshToken = (ExpiringOAuth2RefreshToken) refreshToken;
				Date expiration = expiringRefreshToken.getExpiration();
				if (expiration != null) {
					int seconds = Long.valueOf((expiration.getTime() - System.currentTimeMillis()) / 1000L).intValue();
					pip.expire(refreshKey, seconds);
					pip.expire(refreshAuthKey, seconds);
				}
			}

			pip.sync();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public OAuth2RefreshToken readRefreshToken(String tokenValue) {
		byte[] bytes = conn.get(serializeKey(REFRESH + tokenValue));
		return deserializeRefreshToken(bytes);
	}

	public void removeRefreshToken(OAuth2RefreshToken refreshToken) {
		removeRefreshToken(refreshToken.getValue());
	}

	public void removeRefreshToken(String tokenValue) {
		byte[] refreshKey = serializeKey(REFRESH + tokenValue);
		byte[] refreshAuthKey = serializeKey(REFRESH_TO_AUTH + tokenValue);
		byte[] refresh2AccessKey = serializeKey(REFRESH_TO_ACCESS + tokenValue);
		byte[] access2RefreshKey = serializeKey(ACCESS_TO_REFRESH + tokenValue);

		try (Jedis jedis = conn.getResource()) {
			Pipeline pip = jedis.pipelined();
			pip.del(refreshKey);
			pip.del(refreshAuthKey);
			pip.del(refresh2AccessKey);
			pip.del(access2RefreshKey);
			pip.sync();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public void removeAccessTokenUsingRefreshToken(OAuth2RefreshToken refreshToken) {
		removeAccessTokenUsingRefreshToken(refreshToken.getValue());
	}

	private void removeAccessTokenUsingRefreshToken(String refreshToken) {
		byte[] key = serializeKey(REFRESH_TO_ACCESS + refreshToken);
		List results = null;

		try (Jedis jedis = conn.getResource()) {
			Pipeline pip = jedis.pipelined();
			pip.get(key);
			pip.del(key);
			results = pip.syncAndReturnAll();
		} catch (Exception e) {
			e.printStackTrace();
		}

		if (results != null) {
			byte[] bytes = (byte[]) results.get(0);
			String accessToken = deserializeString(bytes);
			if (accessToken != null) removeAccessToken(accessToken);
		}
	}

	public Collection<OAuth2AccessToken> findTokensByClientIdAndUserName(String clientId, String userName) {
		byte[] approvalKey = serializeKey(UNAME_TO_ACCESS + getApprovalKey(clientId, userName));
		List byteList = null;

		try (Jedis jedis = conn.getResource()) {
			byteList = jedis.lrange(approvalKey, 0L, -1L);
		} catch (Exception e) {
			e.printStackTrace();
		}

		if (byteList != null && byteList.size() != 0) {
			ArrayList accessTokens = new ArrayList(byteList.size());
			for (Object elem : byteList) {
				byte[] bytes = (byte[]) elem;
				OAuth2AccessToken accessToken = deserializeAccessToken(bytes);
				accessTokens.add(accessToken);
			}

			return Collections.unmodifiableCollection(accessTokens);
		} else {
			return Collections.emptySet();
		}
	}

	public Collection<OAuth2AccessToken> findTokensByClientId(String clientId) {
		byte[] key = serializeKey(CLIENT_ID_TO_ACCESS + clientId);
		List byteList = null;

		try (Jedis jedis = conn.getResource()) {
			byteList = jedis.lrange(key, 0L, -1L);
		} catch (Exception e) {
			e.printStackTrace();
		}

		if (byteList != null && byteList.size() != 0) {
			ArrayList accessTokens = new ArrayList(byteList.size());

			for (Object elem : byteList) {
				byte[] bytes = (byte[]) elem;
				OAuth2AccessToken accessToken = deserializeAccessToken(bytes);
				accessTokens.add(accessToken);
			}

			return Collections.unmodifiableCollection(accessTokens);
		} else {
			return Collections.emptySet();
		}
	}

}
